from typing import Optional, Generator

from langchain_openai.chat_models.base import BaseChatOpenAI
from langgraph.types import StreamWriter
from pydantic import BaseModel
from langchain_core.runnables import RunnableConfig

from ...llms import LLMConfig, create_openai_llm
from ...prompts.template import get_prompt_template
from ...utils.image import generate_image_file
from .base import BaseNodes, tool_node
from ...models.graph import GraphState, GraphRequest, AgentStepLog, GraphNodeParam
from loguru import logger


class ImageChatSchema(BaseModel):
    question: str


class ImageGenerateSchema(BaseModel):
    image_desc: str
    sizes: str


class ImageToPageSchema(BaseModel):
    question: str
    path: str


class ImageNodes(BaseNodes):

    def __init__(self, llm_vl: BaseChatOpenAI = None,
                 llm_image: BaseChatOpenAI = None,
                 verbose=False,
                 **kwargs):
        super().__init__(verbose=verbose, **kwargs)
        self.llm_vl = create_openai_llm(llm_vl) if isinstance(llm_vl, LLMConfig) else llm_vl
        self.llm_image = create_openai_llm(llm_image) if isinstance(llm_image, LLMConfig) else llm_image

    @tool_node(args_schema=ImageChatSchema)
    def image_chat(self, state: GraphState, writer: StreamWriter, *, config: Optional[RunnableConfig] = None):
        """
        Description: Have a conversation about an image by asking questions. This tool should be used when you need to describe, or get information about an image through natural language interaction.
        Parameters:
        - question: (required) The question or prompt about the image that you want to ask. Make your question clear and specific.
        Usage:
        <image_chat>
        <question>Your question about the image here</question>
        </image_chat>

        Example:
        <image_chat>
        <question>What objects can you see in this image?</question>
        </image_chat>
        """
        if not self.llm_vl:
            logger.error("Not provided with a vision language model")
            return self.node_failure(message="Not provided with a vision language model")
        try:
            request = GraphRequest.from_runnable_config(config)
            question = state.goto.params.get("question", "describe the image")
            logger.info(f"Processing image_chat with question: '{question}'")

            messages = [{
                "role": "user",
                "content": [
                    {"type": "text", "text": question},
                    {"type": "image_url", "image_url": {"url": f'{get_image_prefix()}{request.image}'}}
                ]
            }]

            last_line = ""
            for chunk in self.llm_vl.stream(messages):
                last_line += chunk.content
                writer(AgentStepLog.build_answer(chunk=last_line, meta={"finish": False}))
            writer(AgentStepLog.build_answer(chunk=last_line, meta={"finish": True}))
            return self.node_success(message=last_line)
        except Exception as e:
            logger.error(f"Error in image chat: {e}")
            return self.node_failure(message="Error in image chat")

    @tool_node(args_schema=ImageGenerateSchema)
    def image_generate(self, state: GraphState, writer: StreamWriter, *, config: Optional[RunnableConfig] = None):
        """
        Description: Generate images based on a text description in multiple sizes. The tool will create detailed, high-quality images matching the provided description and style requirements.
        Parameters:
        - image_desc: (required) A detailed description of the image to generate. Include style, shape, theme, and content details. The description should be specific and avoid colloquial language.
        - sizes: (required) List of image sizes in "widthxheight" format, separated by commas. Example: "512x512,1024x1024"
        Usage:
        <image_generate>
        <image_desc>Your detailed image description here</image_desc>
        <sizes>512x512,1024x1024</sizes>
        </image_generate>

        Example:
        <image_generate>
        <image_desc>A serene mountain landscape at sunset, painted in impressionist style with bold brushstrokes, featuring snow-capped peaks reflected in a crystal-clear alpine lake, warm orange and purple hues in the sky</image_desc>
        <sizes>512x512,1024x1024</sizes>
        </image_generate>
        """
        if not self.llm_image:
            logger.error("Not provided with a vision language model")
            return self.node_failure(message="Not provided with a image generate model")
        image_desc = state.goto.params.get("image_desc", "generate an blank image")
        sizes = state.goto.params.get("sizes", "512x512")

        try:
            image_size = [size.strip() for size in sizes.split(",")]

            images = generate_image_file(
                self.llm_image,
                image_desc,
                image_size,
                "png"
            )

            logger.info(f"Image generation successful: {images}")
            return self.node_success(message=f"Image generated successfully. {images}")
        except Exception as e:
            logger.error(f"Error in image generation: {e}")
            return self.node_failure(message="Error in image generation")

    @tool_node(args_schema=ImageToPageSchema)
    def image_to_page(self, state: GraphState, writer: StreamWriter, *, config: Optional[RunnableConfig] = None):
        """
        Description: Generate a page implementation based on an input image
        Parameters:
        - path: (required) The relative file path where the generated page should be saved (relative to current working directory)
        - question: (required) Additional user question or requirements for the page generation
        Usage:
        <image_to_page>
        <path>the relative path to the generated page</path>
        <question>I want a responsive layout with a dark theme</question>
        </image_to_page>

        Example:
        <image_to_page>
        <path>generated_page.html</path>
        <question>Please generate a landing page with a hero section and contact form</question>
        </image_to_page>
        """
        if not self.llm_vl:
            logger.error("Not provided with a vision language model")
            return self.node_failure(message="Not provided with a vision language model")

        try:
            request = GraphRequest.from_runnable_config(config)
            path = state.goto.params.get("path", "generated_page.html")
            question = state.goto.params.get("question", "generate a html beautiful page")

            if not request.image:
                logger.error("No image provided for image_to_page")
                return self.node_failure(message="No image provided")

            action = "write_to_file"
            last_line = None
            for line in ai_generate_page_by_image(question=question, image=request.image, llm=self.llm_vl):
                if not line:
                    continue
                last_line = line
                writer(AgentStepLog.build_tool_action(action=action,
                                                      output={"type": action,
                                                              "params": {"path": path, "content": last_line}},
                                                      meta={"finish": False}))
            writer(AgentStepLog.build_tool_action(action=action,
                                                  output={"type": action,
                                                          "params": {"path": path, "content": last_line}},
                                                  meta={"finish": True}))
            logger.info(f"Page generation completed successfully, saved to: {path}")
            return self.node_success(message=f"Page generated successfully at {path}, page content:\n\n {last_line}",
                                     force_goto=GraphNodeParam(type=action))
        except Exception as e:
            logger.error(f"Error in image to page: {e}")
            return self.node_failure(message="Error in image to page")


def ai_generate_page_by_image(question: str,
                              image: str,
                              llm: BaseChatOpenAI,
                              verbose=False) -> Generator:
    """
    根据图片生成页面代码
    """

    def _extract_page_content(text: str) -> str:
        if "<page_content>" in text:
            start = text.find("<page_content>") + len("<page_content>")
            if "</page_content>" in text:
                end = text.find("</page_content>")
                return text[start:end].strip()
            return text[start:].strip()
        return ""

    try:
        messages = [{
            "role": "user",
            "content": [
                {"type": "text", "text": get_prompt_template("tool_image2page", question=question)},
                {"type": "image_url", "image_url": {"url": f'{get_image_prefix()}{image}'}}
            ]
        }]

        last_line = ""
        last_content = ""

        for chunk in llm.stream(messages):
            last_line += chunk.content
            if len(last_line) < 30:
                continue

            output = last_line[:-30]
            content = _extract_page_content(output)
            if content and content != last_content:
                yield content
                last_content = content

        if last_line:
            content = _extract_page_content(last_line)
            if content and content != last_content:
                yield content

        logger.info("Page generation from image completed successfully")
    except Exception as e:
        logger.error(f"Failed to ai_generate_page_by_image: {e}")
        raise Exception("Failed to generate page by image")


def get_image_prefix():
    return "data:image/jpeg;base64,"
